#include <iostream>
#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>

static cv::Vec4f _fit_line_2d_weight_least_square(const std::vector<cv::Point2f>& points, const std::vector<float>& weights) {
    int count = points.size();
    CV_Assert(count > 0);
    
    double x = 0, y = 0, x2 = 0, y2 = 0, xy = 0, w = 0;
    if (weights.size() != count) {
        for (int i = 0; i < count; i++) {
            x += points[i].x;
            y += points[i].y;
            x2 += points[i].x * points[i].x;
            y2 += points[i].y * points[i].y;
            xy += points[i].x * points[i].y;
        }
        w = (float)count;
    } else {
        for (int i = 0; i < count; i++) {
            x += weights[i] * points[i].x;
            y += weights[i] * points[i].y;
            x2 += weights[i] * points[i].x * points[i].x;
            y2 += weights[i] * points[i].y * points[i].y;
            xy += weights[i] * points[i].x * points[i].y;
            w += weights[i];
        }
    }
    x /= w;
    y /= w;
    x2 /= w;
    y2 /= w;
    xy /= w;
    float theta = (float)std::atan2(2 * (xy - x * y), (x2 - x * x) - (y2 - y * y)) / 2.;
    cv::Vec4f line;
    line[0] = (float)std::cos(theta);
    line[1] = (float)std::sin(theta);
    line[2] = (float)x;
    line[3] = (float)y;
    return line;
}

static double _calc_distance_2d(const std::vector<cv::Point2f>& points, const cv::Vec4f& line, std::vector<float>& distances) {
    float vx = line[0], vy = line[1], x = line[2], y = line[3];
    distances.resize(points.size());
    double sumDistances = 0;
    for (int i = 0; i < points.size(); i++) {
        float X = points[i].x;
        float Y = points[i].y;
        distances[i] = (float)std::fabs(vx*(Y-y) - vy*(X-x));
        sumDistances += distances[i];
    }
    return sumDistances;
}

static std::vector<float> _calc_weights(int distType, const std::vector<float>& distances, float C=0.f) {
    int count = distances.size();
    std::vector<float> weights(count);
    const double eps = 1e-6;
    for (int i = 0; i < count; i++) {
        switch (distType) {
            case cv::DIST_L1:
                weights[i] = (float)(1.0 / std::max(std::fabs((double)distances[i]), eps));
                break;
            case cv::DIST_L12:
                //weights[i] = 0.5f / (float)(std::sqrt(1. + (double)(distances[i]*distances[i]*0.5)) - 1.);
                weights[i] = 1.f / (float)std::sqrt(1. + (double)(distances[i]*distances[i]*0.5));
                break;
            case cv::DIST_FAIR:
                C = (C == 0 ? 1.3998f : C);
                //weights[i] = 1.f/(C*C*(distances[i]/C-std::log(1+distances[i]/C)));
                weights[i] = 1.f / (1.f+distances[i]/C);
                break;
            case cv::DIST_WELSCH:
                C = (C == 0 ? 2.9846f : C);
                //weights[i] = 1.f/(float)(C*C*0.5*(1.-std::exp(-distances[i]*distances[i]/C/C)));
                weights[i] = (float)std::exp(-distances[i]*distances[i]/(C*C));
                break;
            case cv::DIST_HUBER:
                C = (C <= 0 ? 1.345f : C);
                if (distances[i] < C) {
                    weights[i] = 1.f;
                } else {
                    //weights[i] = 1.f / (C*(distances[i]-C*0.5));
                    weights[i] = C/distances[i];
                }
                break;
            default:
                weights[i] = 1;
                break;
        }
    }
    return weights;
}

static cv::Vec4f _fit_line_2d(const std::vector<cv::Point2f>& points, int distType, float C=0.f, float radiusEps=0.f, float angleEps=0.f, int randomTimes=20, int iter=30) {
    int count = points.size();
    double epsSumDistances = count * FLT_EPSILON;
    double sumDistances = 0, minSumDistances = DBL_MAX;
    cv::Vec4f line, preLine, resultLine;
    resultLine[0] = 0.f;
    resultLine[1] = 0.f;
    resultLine[2] = 0.f;
    resultLine[3] = 0.f;
    float radiusDelta = (radiusEps != 0 ? radiusEps : 1.0f);
    float angleDelta = (angleEps != 0 ? angleEps : 0.01f);
    
    cv::RNG rng((uint64)-1);
    std::vector<float> weights(count);
    for (int k = 0; k < randomTimes; k++) {
        for (int i = 0; i < count; i++) {
            weights[i] = 0.f;
        }
        for (int i = 0; i < std::min(count, 10); ) {
            int j = rng.uniform(0, count);
            if (weights[j] < FLT_EPSILON) {
                weights[j] = 1.f;
                i++;
            }
        }
        line = _fit_line_2d_weight_least_square(points, weights);
        for (int i = 0; i < iter; i++) {
            if (i != 0) {
                double t = preLine[0] * line[0] + preLine[1]*line[1];
                t = std::max(t, -1.);
                t = std::min(t, 1.);
                if (std::fabs(std::acos(t)) < angleDelta) {
                    float x = (float)std::fabs(line[2] - preLine[2]);
                    float y = (float)std::fabs(line[3] - preLine[3]);
                    float d = (x > y ? x : y);
                    if (d < radiusDelta) {
                        break;
                    }
                }
            }
            
            std::vector<float> distances;
            sumDistances = _calc_distance_2d(points, line, distances);
            if (sumDistances < minSumDistances) {
                minSumDistances = sumDistances;
                resultLine[0] = line[0];
                resultLine[1] = line[1];
                resultLine[2] = line[2];
                resultLine[3] = line[3];
                if (sumDistances < epsSumDistances) {
                    break;
                }
            }
            
            weights = _calc_weights(distType, distances, C);
            double sumWeights = 0;
            for (int j = 0; j < count; j++) {
                sumWeights += weights[j];
            }
            if (std::fabs(sumWeights) > FLT_EPSILON) {
                sumWeights = 1. / sumWeights;
                for (int j = 0; j < count; j++) {
                    weights[j] = (float)(weights[j]*sumWeights);
                }
            } else {
                for (int j = 0; j < count; j++) {
                    weights[j] = 1.f;
                }
            }
            
            preLine[0] = line[0];
            preLine[1] = line[1];
            preLine[2] = line[2];
            preLine[3] = line[3];
            line = _fit_line_2d_weight_least_square(points, weights);
        }
        if (sumDistances < minSumDistances) {
            minSumDistances = sumDistances;
            resultLine[0] = line[0];
            resultLine[1] = line[1];
            resultLine[2] = line[2];
            resultLine[3] = line[3];
            if (sumDistances < epsSumDistances) {
                break;
            }
        }
    }
    return resultLine;
}

int main(int argc, char **argv) {
    cv::Mat src = cv::Mat::zeros(600, 600, CV_8UC3);
    std::vector<cv::Point> points = {cv::Point(50, 100),cv::Point(100, 150),cv::Point(150, 200),cv::Point(200, 200),
        cv::Point(200, 250),cv::Point(250, 50),cv::Point(300, 300),cv::Point(300, 350),
        cv::Point(350, 500),cv::Point(400, 400),cv::Point(450, 450)};
    for (int i = 0; i < points.size(); i++) {
        cv::circle(src, points[i], 5, cv::Scalar(255, 255, 255), -1);
    }
    std::vector<cv::String> strs = {"L2", "L1", "L12", "FAIR", "WELSCH", "HUBER"};
    std::vector<cv::Scalar> colors = {cv::Scalar(0, 255, 255),cv::Scalar(255, 0, 255),cv::Scalar(255, 255, 0),
                                    cv::Scalar(255, 0, 0),cv::Scalar(0, 255, 0),cv::Scalar(0, 0, 255)};
    std::vector<int> types = {cv::DIST_L2, cv::DIST_L1, cv::DIST_L12, cv::DIST_FAIR, cv::DIST_WELSCH, cv::DIST_HUBER};
    cv::Mat srcOrigin = src.clone();
    for (int i = 0; i < types.size(); i++) {
        cv::Vec4f line;
        cv::fitLine(points, line, types[i], 0, 0.01, 0.01);
        std::cout << strs[i] << " -- " << line << std::endl;
        float vx = line[0];
        float vy = line[1];
        float x = line[2];
        float y = line[3];
        int lefty = (int)(-x * vy / vx + y);
        int righty = (int)((src.cols-x) * vy / vx + y);
        //cv::Mat m = srcOrigin.clone();
        //cv::line(m, cv::Point(src.cols-1, righty), cv::Point(0, lefty), colors[i], 2);
        //cv::line(src, cv::Point(src.cols-1, righty), cv::Point(0, lefty), colors[i], 2);
        //cv::imshow(strs[i], m);
        
        //if (types[i] == cv::DIST_L1) {
        //if (types[i] == cv::DIST_L12) {
        //if (types[i] == cv::DIST_FAIR) {
        //if (types[i] == cv::DIST_WELSCH) {
        if (types[i] == cv::DIST_HUBER) {
            cv::Mat m = srcOrigin.clone();
            cv::line(m, cv::Point(src.cols-1, righty), cv::Point(0, lefty), colors[i], 1);
            cv::imshow(strs[i], m);
        }
    }
    
    std::vector<cv::Point2f> points_float;
    for (int i = 0; i < points.size(); i++) {
        points_float.push_back(cv::Point2f((float)points[i].x, (float)points[i].y));
    }
    cv::Vec4f line;
    std::vector<float> weights;
    line = _fit_line_2d_weight_least_square(points_float, weights);
    std::cout << strs[0] << " -- my -- " << line << std::endl;
    
    line = _fit_line_2d(points_float, types[1], 0, 0.01, 0.01);
    std::cout << strs[1] << " -- my -- " << line << std::endl;
    
    line = _fit_line_2d(points_float, types[2], 0, 0.01, 0.01);
    std::cout << strs[2] << " -- my -- " << line << std::endl;
    
    line = _fit_line_2d(points_float, types[3], 0, 0.01, 0.01);
    std::cout << strs[3] << " -- my -- " << line << std::endl;
    
    line = _fit_line_2d(points_float, types[4], 0, 0.01, 0.01);
    std::cout << strs[4] << " -- my -- " << line << std::endl;
    
    line = _fit_line_2d(points_float, types[5], 0, 0.01, 0.01);
    std::cout << strs[5] << " -- my -- " << line << std::endl;
    
    float vx = line[0];
    float vy = line[1];
    float x = line[2];
    float y = line[3];
    int lefty = (int)(-x * vy / vx + y);
    int righty = (int)((src.cols-x) * vy / vx + y);
    cv::line(src, cv::Point(src.cols-1, righty), cv::Point(0, lefty), cv::Scalar::all(255), 1);
    
    cv::imshow("src", src);
    cv::waitKey(0);
    return 0;
}
